#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版的Real-ESRGAN微调脚本
避免复杂路径和除零错误
"""

import os
import sys
import yaml
import subprocess
import numpy as np
import torch
# 检查是否有CUDA支持
HAS_CUDA = torch.cuda.is_available()

# 修复torchvision兼容性问题
def fix_torchvision_compatibility():
    """修复torchvision缺少functional_tensor模块的问题"""
    print("正在修复torchvision兼容性问题...")
    try:
        import torchvision.transforms.functional_tensor
        print("torchvision.transforms.functional_tensor模块已存在，无需修复")
    except ImportError:
        print("torchvision.transforms.functional_tensor模块不存在，创建替代模块...")
        # 创建一个简单的替代模块
        class DummyFunctionalTensor:
            def __init__(self):
                pass
            
            def rgb_to_grayscale(self, img, num_output_channels=1):
                """简单的灰度转换实现"""
                import torch
                if num_output_channels == 1:
                    return torch.mean(img, dim=1, keepdim=True)
                else:
                    return torch.cat([torch.mean(img, dim=1, keepdim=True)] * 3, dim=1)
            
            # 添加其他可能需要的方法
            def adjust_contrast(self, img, contrast_factor):
                return img
            
            def adjust_brightness(self, img, brightness_factor):
                return img
        
        # 将替代模块添加到sys.modules
        sys.modules['torchvision.transforms.functional_tensor'] = DummyFunctionalTensor()
        print("torchvision兼容性问题修复完成")

# 先应用兼容性修复
fix_torchvision_compatibility()

# 确保使用正确的Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def create_minimal_config():
    """创建极简配置文件"""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 创建必要的目录
    config_dir = os.path.abspath(os.path.join(current_dir, 'experiments', 'minimal_config'))
    os.makedirs(config_dir, exist_ok=True)
    
    # 使用现有的数据集路径 - 选择REDS_processed数据集
    gt_path = 'D:\\AIJumpProject\\LowLevelVision\\datasets\\REDS_processed\\HR'
    lq_path = 'D:\\AIJumpProject\\LowLevelVision\\datasets\\REDS_processed\\LR'
    print(f"使用数据集: HR={gt_path}, LR={lq_path}")
    
    # 工作目录
    work_dir = os.path.abspath(os.path.join(current_dir, 'experiments', 'minimal_config'))
    os.makedirs(work_dir, exist_ok=True)
    
    # 极简配置 - 避免除零错误的关键配置
    config = {
        'name': 'minimal_realesrgan',
        'work_dir': work_dir,
        'scale': 1,  # 设置为1，因为我们使用的是配对数据
        'manual_seed': 42,
        'mode': 'train',
        'dist': False,
        'auto_resume': False,
        'num_gpu': 0,  # 设置为0表示使用CPU
        'model_type': 'SRModel',
        
        # 数据集配置 - 使用现有的HR/LR图像对
        'datasets': {
            'train': {
                'name': 'LED_Degradation',
                'type': 'PairedImageDataset',
                'dataroot_gt': gt_path,
                'dataroot_lq': lq_path,
                'io_backend': {'type': 'disk'},
                'gt_size': 32,  # 减小图像尺寸以确保匹配
                'use_hflip': False,
                'use_rot': False,
                'batch_size_per_gpu': 1,
                'num_worker_per_gpu': 0,
                'dataset_enlarge_ratio': 1,
                'shuffle': False,
                'scale': 1,
                'phase': 'train'
            }
        },
        
        # 网络结构 - 最简单的配置
        'network_g': {
            'type': 'RRDBNet',
            'num_in_ch': 3,
            'num_out_ch': 3,
            'num_feat': 8,
            'num_block': 1,
            'num_grow_ch': 4
        },
        
        # 训练配置
        'train': {
            'total_iter': 5,
            'val_freq': 5,
            'optim_g': {
                'type': 'Adam',
                'lr': 1e-5,
                'betas': [0.9, 0.99]
            },
            # 损失函数配置
            'pixel_opt': {
                'type': 'L1Loss',
                'loss_weight': 1.0,
                'reduction': 'mean'
            },
            # 添加必要的scheduler配置
            'scheduler': {
                'type': 'MultiStepLR',
                'milestones': [2],
                'gamma': 0.5
            }
        },
        
        # 路径配置
        'path': {
            'experiments_root': work_dir,
            'models': os.path.join(work_dir, 'models'),
            'log': os.path.join(work_dir, 'log')
        },
        
        # 记录器配置
        'logger': {
            'print_freq': 1,
            'save_checkpoint_freq': 5,
            'use_tb_logger': True,
            'wandb': {'project': None, 'resume_id': None}
        }
    }
    
    # 保存配置文件
    config_path = os.path.join(config_dir, 'config.yml')
    with open(config_path, 'w', encoding='utf-8') as f:
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    print(f"配置文件已创建: {config_path}")
    return config_path

def run_training(config_path):
    """运行训练"""
    print(f"\n开始训练，配置文件: {config_path}")
    root_path = os.path.dirname(config_path)
    print(f"根目录路径: {root_path}")
    
    # 创建一个独立的训练脚本文件
    train_script_path = os.path.abspath(os.path.join(os.path.dirname(config_path), 'run_train.py'))
    
    # 脚本内容 - 解析命令行参数并传递root_path
    script_content = '''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""临时训练脚本"""

import os
import sys
import argparse

# 修复torchvision兼容性
class DummyFunctionalTensor:
    def __init__(self):
        pass
    
    def rgb_to_grayscale(self, img, num_output_channels=1):
        import torch
        if num_output_channels == 1:
            return torch.mean(img, dim=1, keepdim=True)
        else:
            return torch.cat([torch.mean(img, dim=1, keepdim=True)] * 3, dim=1)
    
    def adjust_contrast(self, img, contrast_factor):
        return img
    
    def adjust_brightness(self, img, brightness_factor):
        return img

sys.modules['torchvision.transforms.functional_tensor'] = DummyFunctionalTensor()

# 解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='配置文件路径')
args = parser.parse_args()

# 从配置文件路径提取根目录
opt_path = args.opt
root_path = os.path.dirname(os.path.abspath(opt_path))

print(f"配置文件路径: {opt_path}")
print(f"根目录路径: {root_path}")

# 导入并运行训练
from basicsr.train import train_pipeline

# 调用train_pipeline并传递root_path参数
print(f"调用train_pipeline(root_path='{root_path}')")
train_pipeline(root_path)
'''
    
    # 保存脚本文件
    with open(train_script_path, 'w', encoding='utf-8') as f:
        f.write(script_content)
    
    print(f"创建临时训练脚本: {train_script_path}")
    
    # 尝试运行独立脚本
    try:
        print("\n=== 方法3: 运行独立训练脚本 ====")
        # 添加-opt参数指向配置文件
        cmd = [sys.executable, train_script_path, '-opt', config_path]
        print(f"执行命令: {' '.join(cmd)}")
        
        # 强制使用CPU
        torch.cuda.is_available = lambda: False
        
        # 运行脚本并显示实时输出
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding='utf-8',
            errors='replace'
        )
        
        # 实时显示输出
        while True:
            line = process.stdout.readline()
            if not line and process.poll() is not None:
                break
            if line:
                print(line.strip())
        
        return_code = process.wait()
        print(f"训练脚本返回码: {return_code}")
        
        if return_code == 0:
            print("训练成功完成！")
            return True
    except Exception as e:
        print(f"方法3失败: {e}")
        import traceback
        traceback.print_exc()
        print("\n注意：当前环境没有CUDA支持，正在尝试CPU训练")
    
    print("所有训练方法都失败了")
    return False

def main():
    """主函数"""
    print("===== 简化版Real-ESRGAN微调脚本 =====")
    
    # 数据集路径
    gt_path = 'D:\\AIJumpProject\\LowLevelVision\\datasets\\REDS_processed\\HR'
    lq_path = 'D:\\AIJumpProject\\LowLevelVision\\datasets\\REDS_processed\\LR'
    
    # 创建极简配置
    config_path = create_minimal_config()
    
    # 运行训练
    success = run_training(config_path)
    
    if success:
        print("\n🎉 微调训练成功完成！")
        return 0
    else:
        print("\n❌ 微调训练失败")
        print("\n提示：可能是配置问题")
        print(f"数据集路径: HR={gt_path}, LR={lq_path}")
        return 1

if __name__ == '__main__':
    sys.exit(main())